# -*- coding: utf-8 -*-

import os
import sys
import argparse
import binascii

def dargprase():
    parser = argparse.ArgumentParser(description='本程序主要实现 HEX 文件相关操作')
    parser.add_argument('--model', type=str, required=True, help='程序运行模式：1：view, 2: merge')
    parser.add_argument('--apphex', type=str, required=False, help='APP HEX 文件')
    parser.add_argument('--a2lhex', type=str, required=False, help='A2L HEX 文件')
    parser.add_argument('--a2laddr', type=str, required=False, help='A2L 标定文件中 RAM 地址')
    parser.add_argument('--a2llen', type=str, required=False, help='A2L 标定文件中 RAM 长度')
    parser.add_argument('--appaddr', type=str, required=False, help='APP 文件中 ROM 地址')
    args = parser.parse_args()

    if args.model == '1' :
        if not args.apphex:
            parser.error('在 view 模式下，--apphex 参数是必要的，但未提供')
    elif args.model == '2' :
        if not args.apphex:
            parser.error('在 merge 模式下，--apphex 参数是必要的，但未提供')
        if not args.a2lhex:
            parser.error('在 merge 模式下，--a2lhex 参数是必要的，但未提供')
        if not args.a2laddr:
            parser.error('在 merge 模式下，--a2laddr 参数是必要的，但未提供')
        if not args.a2llen:
            parser.error('在 merge 模式下，--a2llen 参数是必要的，但未提供')
        if not args.appaddr:
            parser.error('在 merge 模式下，--appaddr 参数是必要的，但未提供')

    return args

def dcalculate_checksum(data):
    """Calculate the checksum for a line of hex data."""
    return (0x100 - (sum(data) % 256) & 0xFF)

def dread_hex_file(file_path):
    with open(file_path, 'r') as file:
        content = file.readlines()
        return content

def dwrite_hex_file(file_path, content):
    with open(file_path, 'w') as f:
        f.writelines(content)

def dparse_hex_record(record):
    data_length = int(record[1:3], 16)
    record_type = int(record[7:9], 16)
    if record_type == 0x04: # Extended Linear Address Records
        new_base_address = int(record[9:13], 16) << 16
        return record_type, new_base_address, None, data_length
    elif record_type == 0x00:   # Data Records
        address = int(record[3:7], 16)
        data = [int(record[i:i + 2], 16) for i in range(9, 9 + data_length * 2, 2)] # 生成了一个数组
        return record_type, address, data, data_length
    elif record_type == 0x05:   # Start Linear Address Records
        app_address = int(record[9:13], 16) << 16 + int(record[17:21], 16)
        return record_type, app_address, None, data_length
    elif record_type == 0x01:   # End-Of-File Record
        return record_type, None, None, data_length
    else:
        return None, None, None, None

def ddisplay_hex_content(hex_content):
    base_address = 0
    for line in hex_content:
        line = line.strip() # 移除字符串line的前导和尾随空白字符（包括空格、制表符、换行符等）
        record_type, address, data, data_len= dparse_hex_record(line)
        if record_type == 0x04:
            base_address = address
            print(f"Record Type: {record_type}, Address: 0x{address:08X}")
        if data is not None:
            address += base_address
            hex_data = ' '.join(['{:02X}'.format(byte) for byte in data])
            print(f"Record Type: {record_type}, Address: 0x{address:08X}, Data: {hex_data}")

def dprint_hex(hexname):
    try:
        hex_content = dread_hex_file(hexname)
        ddisplay_hex_content(hex_content)
    except FileNotFoundError:
        print(f"File not found: {hexname}")
    except Exception as e:
        print(f"An error occurred: {e}")

def dchange_to_2length(hex_contents):
    new_contents = []
    base_address = 0
    i = 0
    while i < len(hex_contents):
        line1 = hex_contents[i]
        line1 = line1.strip() # 移除字符串line的前导和尾随空白字符（包括空格、制表符、换行符等）
        record_type, address, data, data_len = dparse_hex_record(line1)
        if record_type == 0x04:
            base_address = address
            newline = hex_contents[i]
            i += 1
        elif record_type == 0x00:
            line2 = hex_contents[i+1]
            line2 = line2.strip()
            record_type1, address1, data1, data_len1 = dparse_hex_record(line2)
            if record_type1 == 0x00:
                data = data + data1
                checkdata = [data_len+data_len1, int(line1[3:5], 16), int(line1[5:7], 16), int(line1[7:9], 16)] + data
                checksum = dcalculate_checksum(checkdata)
                newline = f":{data_len+data_len1:02X}{address:04X}{record_type:02X}{''.join(f'{b:02X}' for b in data)}{checksum:02X}\n"
            else:
                newline = hex_contents[i] + hex_contents[i+1]
            i += 2
        else:
            newline = hex_contents[i]
            i += 1
        new_contents.append(newline)

    return new_contents

def dchange_to_2length2(hex_contents):
    new_contents = []
    base_address = 0
    i = 0;
    while i < len(hex_contents):
        line1 = hex_contents[i]
        line1 = line1.strip() # 移除字符串line的前导和尾随空白字符（包括空格、制表符、换行符等）
        record_type, address, data, data_len = dparse_hex_record(line1)
        if record_type == 0x04:
            base_address = address
            newline = hex_contents[i]
            i += 1
        elif record_type == 0x00 and address != 0x0000:
            line2 = hex_contents[i+1]
            line2 = line2.strip()
            record_type1, address1, data1, data_len1 = dparse_hex_record(line2)
            if record_type1 == 0x00:
                data = data + data1 # 数据相加
                checkdata = [data_len+data_len1, int(line1[3:5], 16), int(line1[5:7], 16), int(line1[7:9], 16)] + data
                checksum = dcalculate_checksum(checkdata)
                newline = f":{data_len+data_len1:02X}{address:04X}{record_type:02X}{''.join(f'{b:02X}' for b in data)}{checksum:02X}\n"
            else:
                newline = hex_contents[i] + hex_contents[i+1]
            i += 2
        else:
            newline = hex_contents[i]
            i += 1
        new_contents.append(newline)

    return new_contents

def dchange_line_length(hexname, newname, newlen):
    try:
        hex_contents = dread_hex_file(hexname)
        new_contents = dchange_to_2length(hex_contents)
        dwrite_hex_file(newname, new_contents)
    except FileNotFoundError:
        print(f"File not found: {hexname}")
    except Exception as e:
        print(f"An error occurred: {e}")

def dfind_record_by_addr(contents, addr):
    def_l = 0
    for ln in range(0, len(contents)):
        rt, offset, data, dl = dparse_hex_record(contents[ln])
        if rt == 0x04:
            new_addr = offset
        elif rt == 0x00:
            if def_l < dl:  # 找到行中最大的数据（目前知道的都是16的整数倍）
                def_l = dl
            if new_addr + offset == addr:
                print(f"Find the line[{ln}] by addr[{new_addr + offset:08X}]")
                return new_addr, offset, dl, ln
    
    return None, None, def_l, None

def dfind_position_add_init_line(contents, addr, def_l):
    pre_line_num = 0
    for ln in range(0, len(contents)):
        t, offset, data, dl = dparse_hex_record(contents[ln])
        if t == 0x04:
            if addr & 0xffff0000 < offset:
                print(contents[ln])
                break
        elif t == 0x0:
            pre_line_num = ln

    # 插入必要的数据
    checkdata = [0x02, 0x00, 0x00, 0x04, ((addr >> 24) & 0xff), ((addr >> 16) & 0xff)]
    checksum = dcalculate_checksum(checkdata)
    target_record = f":{''.join(f'{b:02X}' for b in checkdata)}{checksum:02X}\n"
    contents.insert(pre_line_num+1, target_record)  # insert是在当前行前插入的
    print(target_record)
    checkdata = [0, ((addr >> 8) & 0xff), (addr & 0xff), 0x00]
    checksum = dcalculate_checksum(checkdata)
    target_record = f":{''.join(f'{b:02X}' for b in checkdata)}{checksum:02X}\n"
    # 插入新的record, 需要插入无数据之前
    contents.insert(pre_line_num+2, target_record)
    print(target_record)

    return addr & 0xffff0000, (addr >> 16) & 0xffff, def_l, pre_line_num+2, contents

def dmerge_hex_data_prepare(contents, line, count):
    t_data = []
    t_len = 0
    i = 0
    while i < count:
        t, o, data, l = dparse_hex_record(contents[line])  # 需要判断非0x00类型
        if t == 0x04:   # 更新的基地址，但无需拷贝
            line += 1
            continue
        # 暂时无需考虑type不是0x00的类型
        if data is None:
            print(f"The source data is none line:{line}")
            exit(1)
        t_data += data
        t_len += l
        line += 1
        i += 1

    return t_data, t_len, line

def dmerge_hex_data_done(contents, line, b_addr, offset, def_l, src_d, src_l, count):
    i = 0
    while i < count:
        done_l = def_l if def_l < (src_l - def_l * i) else (src_l - def_l * i)
        t, o, data, l = dparse_hex_record(contents[line])
        if t == 0x04 and b_addr == o:   # 这里肯定是hex中间的，且从offset从0开始
            b_addr = o
            offset = 0
            line += 1
            continue
        if data is None:    # 目的位置为空，需要创建Record
            new_dst_addr = b_addr + offset
            if new_dst_addr & 0xffff == 0:  # 说明文件到了block尾，需要额外插入04类型地址段
                print(f"Need Extended Linear Address (0x04) !!!")
                checkdata = [0x02, 0x00, 0x00, 0x04, ((new_dst_addr >> 24) & 0xff), ((new_dst_addr >> 16) & 0xff)]
                print(checkdata)
                checksum = dcalculate_checksum(checkdata)
                target_record = f":{''.join(f'{b:02X}' for b in checkdata)}{checksum:02X}\n"
                contents.insert(line, target_record)
                line += 1

            print(f"We need to add one more record[line:{line} offset:{hex(offset)} len:{done_l}]!!!")
            checkdata = [done_l, ((new_dst_addr >> 8) & 0xff), (new_dst_addr & 0xff), 0x00] + src_d[i*done_l:(done_l*(i+1))]
            checksum = dcalculate_checksum(checkdata)
            target_record = f":{''.join(f'{b:02X}' for b in checkdata)}{checksum:02X}\n"
            # 插入新的record, 需要插入无数据之前
            contents.insert(line, target_record)    # contents[:line - 1] + target_record + contents[line - 1:]
            offset = offset + done_l
        else:
            offset = o
            checkdata = [done_l, int(contents[line][3:5], 16), int(contents[line][5:7], 16), 0x00] + src_d[i*done_l:(done_l*(i+1))]
            checksum = dcalculate_checksum(checkdata)
            contents[line] = f":{''.join(f'{b:02X}' for b in checkdata)}{checksum:02X}\n"
            offset = offset + done_l
        line += 1
        i += 1

    return contents, line, b_addr, offset

def dmerge_hex_data(src_file, dst_file, src_addr, src_len, dst_addr):
    blocks = dparse_hex_blocks(src_file)
    for index, block in enumerate(blocks):
        if src_addr == block['start']:
            if src_len > block['length']:
                src_len = block['length']
                print(f"The len is over {src_len} of the 0x{hex(block['start'])}!")

    # 打开文件
    src_contents = dread_hex_file(src_file)
    dst_contents = dread_hex_file(dst_file)
    
    # 找到源地址所在行，以及基地址
    src_base_addr, src_addr_offset, src_record_default_len, src_line = dfind_record_by_addr(src_contents, src_addr)
    if src_base_addr is None:
        print(f"We do not find the line in the {src_file}")
        exit(1)
    # 找到源地址所在行，以及基地址，如果不存在，则直接插入
    dst_base_addr, dst_addr_offset, dst_record_default_len, dst_line = dfind_record_by_addr(dst_contents, dst_addr)
    if dst_base_addr is None:
        print(f"We do not find the line in dst file, so we will insert")
        dst_base_addr, dst_addr_offset, dst_record_default_len, dst_line, \
            dst_contents = dfind_position_add_init_line(dst_contents, dst_addr, dst_record_default_len)

    if src_record_default_len == dst_record_default_len : # 数据长度一样
        print(f"Both hex have the same length!")
        merge_type = 1
    elif src_record_default_len / dst_record_default_len == 2: # 源数据0x20，目标数据0x10
        print(f"The src is 0x20, but the dst is 0x10")
        merge_type = 2
    elif dst_record_default_len / src_record_default_len == 2: # 源数据0x10，目标数据0x20
        print(f"The src is 0x10, but the dst is 0x20")
        merge_type = 3
    else:
        print(f"The source length {src_record_default_len} or dest length {dst_record_default_len} error")
        exit(1)

    # 从指定地址开始读取数据
    merge_len = 0
    dst_i = 0
    while merge_len < src_len:
        if merge_type == 1:
            src_data, src_dlen, src_line = dmerge_hex_data_prepare(src_contents, src_line, 1)
            dst_contents, dst_line, dst_base_addr, \
            dst_addr_offset = dmerge_hex_data_done(
                    dst_contents, dst_line, dst_base_addr, dst_addr_offset, dst_record_default_len,
                    src_data, src_dlen, 1)
            merge_len += src_dlen

        elif merge_type == 2:
            src_data, src_dlen, src_line = dmerge_hex_data_prepare(src_contents, src_line, 1)
            dst_contents, dst_line, dst_base_addr, \
            dst_addr_offset = dmerge_hex_data_done(
                    dst_contents, dst_line, dst_base_addr, dst_addr_offset, dst_record_default_len,
                    src_data, src_dlen, 2)
            merge_len += src_dlen
        elif merge_type == 3:
            src_data, src_dlen, src_line = dmerge_hex_data_prepare(src_contents, src_line, 2)
            dst_contents, dst_line, dst_base_addr, \
            dst_addr_offset = dmerge_hex_data_done(
                    dst_contents, dst_line, dst_base_addr, dst_addr_offset, dst_record_default_len,
                    src_data, src_dlen, 1)
            merge_len += src_dlen
        else:
            print(f"The type not currently supported!!!")
            exit(1)

    # 将数据写入目标hex文件指定地址
    dwrite_hex_file(dst_file, dst_contents)

def dis_valid_hex_file(file):
    try:
        with open(file, 'r') as f:
            for line in f:
                line = line.strip()
                if not line.startswith(':'):
                    print(f"The line start error: {line}")
                    return False
                l = int(line[1:3], 16)
                checksum = int(line[-2:], 16)
                checkdata = [int(line[i:i + 2], 16) for i in range(1, l*2 + 9, 2)]
                if checksum != dcalculate_checksum(checkdata):
                    print(f"The line checksum error: {line}")
                    return False
    except Exception as e:
        print(f"Error reading file: {e}")
        return False
    return True

def dparse_hex_blocks(file):
    blocks = []
    contents = dread_hex_file(file)
    current_block = {'start':0, 'end':0, 'length':0}
    max_len = 0
    last_o = 0
    for cont in contents:
        t, o, data, l = dparse_hex_record(cont)
        if t == 0x04:
            last_baddr = o
            if current_block['start'] == 0: # 第一个0x04类型
                current_block['start'] = o
                current_block['end'] = o
            elif current_block['end'] != o:    # 地址不连续，新一个block
                current_block['length'] = current_block['end'] - current_block['start']
                current_block['end'] -= 1   # 当前block的最后一个字节地址
                blocks.append(current_block.copy())
                # 因为当前已经是新block的0x04类型行，所以需要将偏移量赋值给初始值
                current_block = {'start':o, 'end':o, 'length':0}
                last_o = 0
        elif t == 0x00:
            if l > max_len:
                max_len = l

            if current_block['length'] == 0:  # 第一个可能不是从0开始的
                current_block['start'] += o
                current_block['end'] += o
            else:
                if o > (last_o + max_len):  # 在同一个基地址下，存在不同的block
                    # print(f"{o:X} {last_o:X} {max_len:X}")
                    current_block['length'] = current_block['end'] - current_block['start']
                    current_block['end'] -= 1
                    blocks.append(current_block.copy())
                    last_o = o
                    current_block = {'start':last_baddr + o, 'end':last_baddr + o + l, 'length':l}
                    continue

            current_block['end'] += l
            current_block['length'] += l
            last_o = o
        elif t == 0x01:   # 当前block结束了
            current_block['length'] = current_block['end'] - current_block['start']
            current_block['end'] -= 1
            blocks.append(current_block.copy())
            last_o = 0
            current_block = {'start':0, 'end':0, 'length':0}
    return blocks

def dcheck_and_parse_blocks(file):
    if dis_valid_hex_file(file) is True:
        blocks = dparse_hex_blocks(file)
        print(f"The file is: {file}")
        print(f"File Hex Blocks Information:")
        for index, block in enumerate(blocks):
            print(f"Block\t{index}\tStarts at: 0x{block['start']:08X}\tEnds at: 0x{block['end']:08X}\t(Length: 0x{block['length']:0X}={block['length']})")
        return blocks
    else:
        print(f"{file} not a valid hex file!!!")
        exit(1)

def dremove_hex_blocks(file, block_addr):
    blocks = dparse_hex_blocks(file)
    for idx, block in enumerate(blocks):
        if block_addr == block['start']:
            print(f"Find the blocks addr:{block['start']:08X} len:{block['length']:0X}")
            new_contents = []
            contents = dread_hex_file(file)
            start_i = 0
            end_i = 0
            save_last_04 = 0
            for index, cont in enumerate(contents):
                t, o, data, l = dparse_hex_record(cont)
                if t == 0x04:
                    last_addr = o
                    last_04_cont = cont
                    t, o, data, l = dparse_hex_record(contents[index + 1])
                    if (last_addr + o) == block_addr:
                        print(f"record base addr:{last_addr + o:X}")
                        start_i = index
                elif t == 0x0:
                    if start_i == 0 and (last_addr + o) == block_addr:
                        print(f"record first addr:{last_addr + o:X}")
                        start_i = index
                    elif (last_addr + o) < (block_addr + block['length']) and (last_addr + o) > block_addr:
                        end_i = index

                if (last_addr + o + l) == (block_addr + block['length']):   # 结束了
                    t, o, data, l = dparse_hex_record(contents[index + 1])
                    if t == 0x00:
                        save_last_04 = 1
                    break

            if 1 == save_last_04:
                new_contents.extend(contents[:start_i])
                new_contents.append(last_04_cont)
                new_contents.extend(contents[end_i+1:])
            else:
                new_contents.extend(contents[:start_i])
                new_contents.extend(contents[end_i+1:])
            dwrite_hex_file(file, new_contents)
            return

if __name__ == '__main__':
    # 从命令行参数获取源文件、目标文件、源文件地址、源文件长度、目标文件地址
    margs = dargprase()

    if margs.model == '1':
        print(f"The model is view")
        dcheck_and_parse_blocks(margs.apphex)
        # dchange_line_length(margs.apphex, "aaa.hex", 2)
        # dprint_hex(margs.apphex)
        # dremove_hex_blocks(margs.apphex, 0x20000000)
    elif margs.model == '2':
        print(f"The model is merge")
        dst_file = margs.apphex
        src_file = margs.a2lhex
        dcheck_and_parse_blocks(margs.apphex)
        dcheck_and_parse_blocks(margs.a2lhex)
        src_addr = int(margs.a2laddr, 16)  # 源文件地址，假设为16进制字符串
        src_len = int(margs.a2llen, 16)  # 数据长度，假设为16进制字符串
        dst_addr = int(margs.appaddr, 16)  # 目标文件地址，假设为16进制字符串
        # 调用函数进行数据复制操作
        dmerge_hex_data(src_file, dst_file, src_addr, src_len, dst_addr)
    else:
        print("model 输入异常")

